import sys

# sys.path.insert(0, "../../python")
sys.path.insert(0, '/home/xiaomin/wxm/mxnet/python')
import cv2 as cv
import numpy as np
import sys
import ConfigParser
import os
import mxnet as mx
from scipy.io import loadmat
from sklearn.metrics import roc_auc_score
import datetime
from scikits import audiolab
from matplotlib.pyplot import specgram
from scipy.signal.spectral import spectrogram


class Submitter:
    def __init__(self, config):
        self.config = config
        # model
        self.model_prefix = config.get('model', 'model_prefix')
        self.model_epoch = int(config.get('model', 'model_epoch'))
        self.checkpoint = mx.model.FeedForward.load(self.model_prefix, self.model_epoch)
        self.ctx = int(config.get('mode', 'ctx'))
        self.model = mx.model.FeedForward(symbol=self.checkpoint.symbol,
                                          arg_params=self.checkpoint.arg_params,
                                          aux_params=self.checkpoint.aux_params,
                                          ctx=mx.gpu(self.ctx))
        self.result_dir = config.get('model', 'result_dir')
        self.file_name = config.get('mode', 'file_name')
        if not os.path.isdir(self.result_dir):
            os.mkdir(self.result_dir)

        # data
        self.num_label = int(config.get('data', 'num_label'))
        self.list = config.get('data', 'list')
        self.dir = config.get('data', 'dir')
        self.mode = config.get('mode', 'mode')

    def run_single(self, line):
        line = line
        label = None
        fn = None
        path = None
        if self.mode == 'eval':
            index, path, label = line.strip('\n').split('\t')
        elif self.mode == 'submit':
            index, path = line.strip('\n').split('\t')
            _, fn = path.split('/')

        fn = fn[0:-7] + '.mat'
        print os.path.join(self.dir, path)
        wav = audiolab.wavread(os.path.join(self.dir, path))
        wav_arr = wav[0]
        # Pxx, freqs, bins, im = specgram(wav_arr, NFFT=512, Fs=400, noverlap=154, scale='dB')
        freqs, t, Pxx = spectrogram(wav_arr, nperseg=512, fs=400, noverlap=154, window='hanning', mode='psd')
        Pxx = np.expand_dims(Pxx, 0)
        return Pxx, fn
        # Pxx = np.expand_dims(Pxx, 0)
        # pred = self.model.predict(Pxx).squeeze()
        # if self.mode == 'eval':
        #     return pred, label
        # elif self.mode == 'submit':
        #     return pred, fn

    def run_all(self):
        print 'run all'
        lines = open(self.list, 'r').readlines()
        num = len(open(self.list, 'r').readlines())
        idx = 1
        preds = np.array([])
        fns = []
        Pxxs = []
        for i in range(num):
            line = lines[i]
            print self.model_prefix.split('/')[-1] + '_' + str(self.model_epoch)
            print 'Process ' + str(idx) + 'th ' + 'in ' + str(num)
            Pxx, fn = self.run_single(line)
            print Pxx.shape
            Pxxs.append(Pxx)
            pred = None
            if (i + 1) % 10 == 0:
                Pxxs_arr = np.asarray(Pxxs)
                print Pxxs_arr.shape
                x = mx.io.NDArrayIter(data=Pxxs_arr, label=None, batch_size=10, shuffle=False, last_batch_handle='pad')
                print x
                pred = self.model.predict(x)
                print pred.shape
                Pxxs = []
                if preds is not None:
                    preds = np.concatenate((preds, pred), 0)
                elif preds is None:
                    preds = pred
            fns.append(fn)
            print '*' * 80
            idx += 1

        preds = np.array(preds)
        now = datetime.datetime.now()
        csv_name = str(now.strftime('%Y-%m-%d-%H-%M')) + '_' + self.file_name + '.csv'
        f = open(self.result_dir + csv_name, 'w')
        f.write('File,Class\n')
        print 'writing submission ...'
        for i in range(len(preds)):
            line = fns[i] + ',' + str(preds[i][1]) + '\n'
            f.write(line)


if __name__ == '__main__':
    config_path = 'cfgs/submit_mxnet_cnn.cfg'
    config = ConfigParser.RawConfigParser()
    config.read(config_path)
    miner = Submitter(config)
    miner.run_all()
